package org.example.state;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.example.data.WaterSensor;
import org.example.function.WaterSensorMapFunction;

import java.time.Duration;

public class OperatorListState {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        DataStreamSource<String> source = env.socketTextStream("localhost", 9999);
        source.map(new MyCountMapFunction()).print();


        env.execute("OperatorListState");
    }

    /**
     * 实现CheckPointFunction接口
     */
    public static class MyCountMapFunction implements MapFunction<String, Long>, CheckpointedFunction {
        private Long count = 0L;
        private ListState<Long> state;

        @Override
        public Long map(String value) throws Exception {
            return ++count;
        }

        /**
         * 定义将本地变量的持久化，将 本地变量 拷贝 到算子状态中
         *
         * @param context the context for drawing a snapshot of the operator
         * @throws Exception
         */
        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            System.out.println("调用方法 = snapshotState");

            //清空算子状态
            state.clear();
            //将本地变量添加到算子状态中
            state.add(count);

        }

        /**
         * 初始化本地变量，程序恢复时，从状态中 把数据添加到本地变量，每个子任务执行一次。
         *
         * @param context the context for initializing the operator
         * @throws Exception
         */
        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            System.out.println("调用方法 = initializeState");
            //初始化算子状态
            state = context.getOperatorStateStore().getListState(new ListStateDescriptor<Long>("state", Types.LONG));
            //是否恢复成功
            if (context.isRestored()) {
                Iterable<Long> longs = state.get();
                longs.forEach(c -> {
                    count = count + c;
                });
            }
        }
    }
}
